package main

import (
	"context"
	"fmt"
	"net"
	"os"
	"runtime"
	"shell/common"
	"shell/common/packages"
	"shell/global"
	"shell/proxy/config"
	"shell/utils"
	"time"

	"github.com/dnephin/pflag"
	"github.com/google/uuid"
	"github.com/quic-go/quic-go"
)

var (
	u = pflag.BoolP("uuid", "u", false, "generate a uuid and exit")

	devMode = pflag.BoolP("dev-mode", "d", false, "develop mode, print goroutine number every 10 seconds")

	cfgPath = pflag.StringP("config", "c", "./config.yaml", "config file path")

	globalCfg *config.ProxyConfig = nil

	bytePool = utils.NewBytesPool()

	logger = global.DefaultLogger()
)

func DisplayGoroutine(ctx context.Context) {
	logger.SetLevel(global.Debug)
	ticker := time.NewTicker(time.Second * 10)
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			logger.Debugf(0, "goroutine num: %d", runtime.NumGoroutine())
		}
	}
}

func main() {
	pflag.Parse()

	if *u {
		fmt.Printf("uuid: %s \n", uuid.New().String())
		os.Exit(0)
	}

	cfg, err := config.ParseCfg(*cfgPath)
	if err != nil {
		logger.Fatalf(0, "error parse config: %v", err)
	}
	globalCfg = cfg

	quicListener, err := quic.ListenAddr(net.JoinHostPort("0.0.0.0", fmt.Sprintf("%d", globalCfg.QuicPort)), globalCfg.TLSCfg, &quic.Config{
		MaxIdleTimeout:     15 * time.Second,
		MaxIncomingStreams: 1 << 60,
	})
	if err != nil {
		logger.Fatalf(0, "error listen quic port: %v", err)
	}

	// 将本节点信息写入到全局map中
	now := time.Now()
	stat := new(packages.ControlledStat)
	stat.UUID = globalCfg.UUID
	stat.PrivateIP = net.IPv4(127, 0, 0, 1)
	stat.PublicIP = net.IPv4(0, 0, 0, 0)
	stat.PublicPort = globalCfg.QuicPort
	stat.IsDelay = false
	stat.UseDelay = false
	stat.OperatingSystem = packages.GetOS()
	stat.LastKeepTime = now
	cei, _ := NewControlledEndInfo(stat, nil)
	ControlledEndMap.Store(globalCfg.UUID, cei)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	if *devMode {
		go DisplayGoroutine(ctx)
	}
	timeoutWatcher = common.NewTimeoutWatcher(ctx, 30*time.Second)

	// 处理quic连接请求
	for {
		quicConn, err := quicListener.Accept(ctx)
		if err != nil {
			logger.Errorf(0, "error accept quic connection: %v", err)
			break
		}
		conn := common.NewQuicConn(quicConn, bytePool)
		go handleQuicConn(conn)
	}
}

// 处理quic连接
func handleQuicConn(conn *common.QuicConn) {
	stream, err := conn.AcceptStream()
	if err != nil {
		logger.Errorf(0, "error accept stream: %v", err)
		return
	}
	metadata, err := common.DetectPackage(stream)
	if err != nil {
		common.HandleErr(err)
		stream.CloseConn(common.CloseByProcess)
		return
	}
	role := metadata.Role
	switch role {
	case packages.RoleClient:
		// 是客户端请求
		handleClientEnd(stream, metadata)
	case packages.RoleControlled:
		// 是被控端请求
		handleControlledEnd(stream, metadata)
	default:
		logger.Errorf(0, "error accpet error role !!!")
		stream.CloseConn(common.CloseByProcess)
	}
}
